import numpy as np 
# from test_least_square import least_sq
import torch 
import matplotlib.pyplot as plt 

def batch_leastsq(x, y):
    ## 批量线性最小二乘法计算
    x_bar = torch.mean(x, dim=1, keepdim=True)
    y_bar = torch.mean(y, dim=1, keepdim=True)
    k = ((x - x_bar) * (y - y_bar)).sum(dim=1, keepdim=True) / torch.sum((x - x_bar)**2, dim=1, keepdim=True)
    b = y_bar - k * x_bar
    return k, b 

def extract_k_b(data, size=20):
    # 提取k b作为特征
    # 1. 分段 分20段
    time_step = data.shape[2]
    batch_size = data.shape[0]
    frequency = int(time_step / size)
    channel = data.shape[1]
    k_features = torch.empty(0)
    b_features = torch.empty(0)
    for i in range(size):
        tmp_data = data[:, :, i*frequency:(i+1)*frequency]
        tmp_x = torch.arange(i*frequency+1, (i + 1)*frequency + 1, dtype=torch.float32)
        tmp_x = tmp_x.repeat(batch_size, 1)
       
        channel_k = torch.empty(0)
        channel_b = torch.empty(0)
        for c in range(channel):
            c_temp_data = tmp_data[:, c, :]
        
            k, b = batch_leastsq(tmp_x, c_temp_data)  # 得到一个通道的k，b参数数据
            channel_k = torch.cat((channel_k, k), dim=1)
            channel_b = torch.cat((channel_b, b), dim=1)
           
        if i == 0 :
            k_features = channel_k.unsqueeze(-1)
            b_features = channel_b.unsqueeze(-1)
        else :
            k_features = torch.cat((k_features, channel_k.unsqueeze(-1)), dim=-1)
            b_features = torch.cat((b_features, channel_b.unsqueeze(-1)), dim=-1)

    return k_features, b_features

if __name__ == "__main__":
    # x_i = torch.tensor([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]], dtype=torch.float32) 
    # y_i = torch.tensor([[2, 4, 6.5, 7, 11], [2, 4, 6.5, 7, 11]])
    # print(leastsq(x_i, y_i))
   
    # src_data = torch.load("./src_data_raw.pkl")
    # src_data = src_data[:100]
    # print(src_data.shape)
    # src_data = src_data.permute(0, 2, 1) #(batch, channel, time)
    # src_data = (src_data - torch.mean(src_data)) / torch.std(src_data)

    # k, b = extract_k_b(src_data)
    # print(k.shape)
    # print(b.shape)

    # k_b_features = torch.cat((k, b), dim=1)
    # print(k_b_features.shape)
    # torch.save(k_b_features, "./k_b_features.pkl")

    # src_data = torch.load("./src_data_0628.pkl")
    # # src_data = src_data[:100]
    # print(src_data.shape)
    # src_data = src_data.permute(0, 2, 1)
    # k, b = extract_k_b(src_data, size=15)
    # print(k.shape)
    # print(b.shape)
    # k_b_features = torch.cat((k, b), dim=1)
    # print(k_b_features.shape)
    # torch.save(k_b_features, "./k_b_features0708.pkl")

    test_src_data = torch.load("./src_data_0711.pkl")
    # test_src_data = test_src_data[:100]
    print(test_src_data.shape)
    test_src_data = test_src_data.permute(0, 2, 1)
    k, b = extract_k_b(test_src_data, size=15)
    # print(k)
    # print(b)

    # k_b_features = torch.cat((k, b), dim=1)
    # print(k_b_features.shape)
    # torch.save(k_b_features, "./k_b_features0711.pkl")

    
    k = k[1, 1, :]
    b = b[1, 1, :]

    data = test_src_data[1, 1]
    print(data.shape)
    x = torch.arange(1, 376)
    plt.plot(x, data)
    for i in range(15):
        x = torch.arange(i*25+1, (i+1)*25+1)
        y = k[i]*x + b[i]
        plt.plot(x, y, color="red")
    plt.show()

    # print(k)
    # print(b)    






